{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import git\n",
    "import pandas as pd\n",
    "from fake_useragent import UserAgent\n",
    "import requests\n",
    "import datetime\n",
    "import json\n",
    "from utils.file_utils import dir_check, safe_open\n",
    "from utils.parser_utils import get_query, get_syntax_tree\n",
    "from tree_sitter import Tree\n",
    "from utils.sample_utils import (\n",
    "    find_parent_fn_node,\n",
    "    node_of_unsafe_func,\n",
    "    node_of_pub_func,\n",
    ")\n",
    "crates_info_file = \"data/crate_meta/crates_info.csv\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Crawl crates info from crates.io"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the crates infromation from crates.io\n",
    "def craw_crates_list(crates_info_file):\n",
    "    df = pd.DataFrame(\n",
    "        columns=[\n",
    "            \"name\",\n",
    "            \"newest_version\",\n",
    "            \"recent_downloads\",\n",
    "            \"created_at\",\n",
    "            \"updated_at\",\n",
    "            \"repository\",\n",
    "        ]\n",
    "    )\n",
    "    headers = {\n",
    "        \"User-Agent\": UserAgent().random,\n",
    "        \"referer\": \"https://crates.io/crates?sort=downloads\",\n",
    "    }\n",
    "    total = 0\n",
    "    n = 0\n",
    "    error_num = 0\n",
    "    while 1:\n",
    "        url = (\n",
    "            \"https://crates.io/api/v1/crates?page=\"\n",
    "            + str(n + 1)\n",
    "            + \"&per_page=100&sort=downloads\"\n",
    "        )\n",
    "        try:\n",
    "            page_text = requests.get(url=url, headers=headers).text\n",
    "            crates = json.loads(page_text)[\"crates\"]\n",
    "        except:\n",
    "            print(\"Error in {}, retry-{} \".format(n + 1, error_num))\n",
    "            error_num += 1\n",
    "            if error_num < 10:\n",
    "                continue\n",
    "            else:\n",
    "                break\n",
    "        data_lists = [[] for i in range(6)]\n",
    "        for crate in crates:\n",
    "            data_lists[0].append(crate[\"name\"])\n",
    "            data_lists[1].append(crate[\"newest_version\"])\n",
    "            data_lists[2].append(crate[\"recent_downloads\"])\n",
    "            data_lists[3].append(crate[\"created_at\"])\n",
    "            data_lists[4].append(crate[\"updated_at\"])\n",
    "            data_lists[5].append(crate[\"repository\"])\n",
    "        dic = {\n",
    "            \"name\": data_lists[0],\n",
    "            \"newest_version\": data_lists[1],\n",
    "            \"recent_downloads\": data_lists[2],\n",
    "            \"created_at\": data_lists[3],\n",
    "            \"updated_at\": data_lists[4],\n",
    "            \"repository\": data_lists[5],\n",
    "        }\n",
    "        df2 = pd.DataFrame(dic)\n",
    "        df = pd.concat([df, df2], axis=0)\n",
    "        print(\"Crawl page {}, get {} info\".format(n + 1, len(df2)))\n",
    "        error_num = 0\n",
    "        total += len(df2)\n",
    "        n += 1\n",
    "        if len(df2) < 50:\n",
    "            print(\"Get {} crates in total\".format(total))\n",
    "            break\n",
    "\n",
    "    # crates are sorted with \"recent_downloads\" descendingly\n",
    "    df.to_csv(crates_info_file, index=None)\n",
    "\n",
    "\n",
    "craw_crates_list(crates_info_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter the crates (top-n)\n",
    "def filter_crates(crates_info_file, result_file, target_range=(0,500)):\n",
    "    data = pd.read_csv(crates_info_file, sep=\",\", header=0, usecols=[0, 2, 4, 5])\n",
    "    array = data.values[0::, 0::]\n",
    "    repo_list = []\n",
    "    error_count = [0, 0, 0, 0]\n",
    "    # Rust 2018: 2018-12-07\n",
    "    RUST_2018 = datetime.datetime(2018, 12, 7, 0, 0, 0)\n",
    "    for i in range(target_range[0], target_range[1]):\n",
    "        update = datetime.datetime.strptime(str(array[i][2])[:19], \"%Y-%m-%dT%H:%M:%S\")\n",
    "        if update <= RUST_2018:\n",
    "            error_count[0] += 1\n",
    "            continue\n",
    "\n",
    "        if len(str(array[i][3])) < 5:\n",
    "            error_count[1] += 1\n",
    "            continue\n",
    "        url = str(array[i][3]).split(\"/\")\n",
    "        if \"github.com\" in url[2]:\n",
    "            repo = url[3] + \"/\" + url[4].replace(\".git\", \"\")\n",
    "            if repo not in repo_list:\n",
    "                repo_list.append(repo)\n",
    "            else:\n",
    "                error_count[3] += 1\n",
    "        else:\n",
    "            error_count[2] += 1\n",
    "\n",
    "    print(len(repo_list))\n",
    "    print(error_count)\n",
    "    open(result_file, \"w\").write(\"\\n\".join(repo_list))\n",
    "    return repo_list\n",
    "\n",
    "\n",
    "\n",
    "def clone_repo(repo_path, repo_dir):\n",
    "    # clone repo from github\n",
    "    repo_path = repo_path.strip()\n",
    "    fail_count = 0\n",
    "    while 1:\n",
    "        try:\n",
    "            repo_url = \"https://github.com/\" + repo_path + \".git\"\n",
    "            if os.path.exists(repo_dir):\n",
    "                break\n",
    "            git.Repo.clone_from(url=repo_url, to_path=repo_dir, depth=1)\n",
    "            break\n",
    "        except:\n",
    "            if fail_count < 10:\n",
    "                fail_count += 1\n",
    "                continue\n",
    "            else:\n",
    "                print(repo_path + \" clone failed!\")\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "unsafe_block_pattern = \"(unsafe_block) @unsafe_block\"\n",
    "query = get_query(unsafe_block_pattern)\n",
    "\n",
    "def find_risky_func(repository_dir: str):\n",
    "    samples = list()\n",
    "    for dir_path, _, files in os.walk(repository_dir):\n",
    "        for file_name in files:\n",
    "            if file_name[-3:] != \".rs\":\n",
    "                continue\n",
    "            file_path = os.path.join(dir_path, file_name)\n",
    "            relative_file = os.path.relpath(os.path.join(dir_path, file_name), repository_dir)\n",
    "            if \"test\" in file_path or \"example\" in file_path or \"bench\" in file_path:\n",
    "                continue\n",
    "            tree: Tree = get_syntax_tree(file_path)\n",
    "            captures = query.captures(tree.root_node)\n",
    "            for node, _ in captures:\n",
    "                # tranverse cursor to father function\n",
    "                fn_node = find_parent_fn_node(node)\n",
    "                if fn_node is None:\n",
    "                    continue\n",
    "\n",
    "                # filter unsafe block in unsafe function\n",
    "                if node_of_unsafe_func(fn_node):\n",
    "                    continue\n",
    "\n",
    "                # filter private functions\n",
    "                if not node_of_pub_func(fn_node):\n",
    "                    continue\n",
    "\n",
    "                duplicated = False\n",
    "                for s in samples:\n",
    "                    if s[\"relative_file\"] == relative_file and s[\"start_line\"] == fn_node.start_point[0]:\n",
    "                        duplicated = True\n",
    "                        break\n",
    "                if duplicated:\n",
    "                    continue\n",
    "\n",
    "                samples.append(\n",
    "                    {\n",
    "                        \"relative_file\": relative_file,\n",
    "                        \"start_line\": fn_node.start_point[0], \n",
    "                        \"start_byte\": fn_node.start_byte,\n",
    "                        \"end_line\": fn_node.end_point[0],\n",
    "                        \"end_byte\": fn_node.end_byte,\n",
    "                    }\n",
    "                )\n",
    "    print(f\"Got {len(samples)} risky function in {repository_dir}\")\n",
    "    return samples\n",
    "\n",
    "\n",
    "def find_unsafe_func(repository_dir: str):\n",
    "    samples = list()\n",
    "    for dir_path, _, files in os.walk(repository_dir):\n",
    "        for file_name in files:\n",
    "            if file_name[-3:] != \".rs\":\n",
    "                continue\n",
    "            file_path = os.path.join(dir_path, file_name)\n",
    "            relative_file = os.path.relpath(os.path.join(dir_path, file_name), repository_dir)\n",
    "            if \"test\" in file_path or \"example\" in file_path or \"bench\" in file_path:\n",
    "                continue\n",
    "            tree: Tree = get_syntax_tree(file_path)\n",
    "            captures = query.captures(tree.root_node)\n",
    "            for node, _ in captures:\n",
    "                # tranverse cursor to father function\n",
    "                fn_node = find_parent_fn_node(node)\n",
    "                if fn_node is None:\n",
    "                    continue\n",
    "\n",
    "                # filter unsafe block in unsafe function\n",
    "                if not node_of_unsafe_func(fn_node):\n",
    "                    continue\n",
    "\n",
    "                # filter private functions\n",
    "                if not node_of_pub_func(fn_node):\n",
    "                    continue\n",
    "\n",
    "                duplicated = False\n",
    "                for s in samples:\n",
    "                    if s[\"relative_file\"] == relative_file and s[\"start_line\"] == fn_node.start_point[0]:\n",
    "                        duplicated = True\n",
    "                        break\n",
    "                if duplicated:\n",
    "                    continue\n",
    "        \n",
    "                # filter extern C functions\n",
    "                if \"extern \\\"C\\\"\" in fn_node.text.decode(\"utf-8\"):\n",
    "                    continue\n",
    "\n",
    "                samples.append(\n",
    "                    {\n",
    "                        \"relative_file\": relative_file,\n",
    "                        \"start_line\": fn_node.start_point[0], \n",
    "                        \"start_byte\": fn_node.start_byte,\n",
    "                        \"end_line\": fn_node.end_point[0],\n",
    "                        \"end_byte\": fn_node.end_byte,\n",
    "                    }\n",
    "                )\n",
    "    print(f\"Got {len(samples)} unsafe functions in {repository_dir}\")\n",
    "    return samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get top-500 crates\n",
    "### Clone target repositories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_crate_file = \"data/crate_meta/top_crates.txt\"\n",
    "crates_repo_dir = \"data/crates_repo\"\n",
    "risky_func_dir = \"data/risky_func\"\n",
    "unsafe_func_dir = \"data/unsafe_func\"\n",
    "dir_check(crates_repo_dir)\n",
    "dir_check(risky_func_dir)\n",
    "dir_check(unsafe_func_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "356\n",
      "[6, 3, 5, 130]\n"
     ]
    }
   ],
   "source": [
    "# filter the crates (top-n)\n",
    "top_repos = filter_crates(crates_info_file, top_crate_file, (0, 500))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for repo_path in top_repos:\n",
    "    repo_dir = os.path.join(crates_repo_dir, \".\".join(repo_path.split(\"/\")))\n",
    "    clone_repo(repo_path, repo_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get risky and unsafe samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for repo_path in top_repos:\n",
    "    repo_name = \".\".join(repo_path.split(\"/\"))\n",
    "    repo_dir = os.path.join(crates_repo_dir, repo_name)\n",
    "    unsafe_func = find_unsafe_func(repo_dir)\n",
    "    json.dump(unsafe_func, safe_open(os.path.join(unsafe_func_dir, repo_name + \".json\"), \"w\"), indent=2)\n",
    "    risky_func = find_risky_func(repo_dir)\n",
    "    json.dump(risky_func, safe_open(os.path.join(risky_func_dir, repo_name + \".json\"), \"w\"), indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scan top 500 - 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "scan_crate_file = \"data/crate_meta/scan_crates.txt\"\n",
    "scan_repo_dir = \"data/scan_repo\"\n",
    "scan_func_dir = \"data/scan_func\"\n",
    "dir_check(scan_repo_dir)\n",
    "dir_check(scan_func_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "878\n",
      "[78, 29, 31, 484]\n"
     ]
    }
   ],
   "source": [
    "# filter the crates (top-n)\n",
    "scan_target = filter_crates(crates_info_file, scan_crate_file, (500, 2000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "827"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scan_target = [t for t in scan_target if t not in top_repos]\n",
    "len(scan_target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wasmerio/wasmer clone failed!\n",
      "jmap-rs/json-pointer clone failed!\n",
      "paritytech/pariry-common clone failed!\n"
     ]
    }
   ],
   "source": [
    "for repo_path in scan_target:\n",
    "    repo_dir = os.path.join(scan_repo_dir, \".\".join(repo_path.split(\"/\")))\n",
    "    clone_repo(repo_path, repo_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get risky samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for repo_path in scan_target:\n",
    "    repo_name = \".\".join(repo_path.split(\"/\"))\n",
    "    repo_dir = os.path.join(scan_repo_dir, repo_name)\n",
    "    risky_func = find_risky_func(repo_dir)\n",
    "    json.dump(risky_func, safe_open(os.path.join(scan_func_dir, repo_name + \".json\"), \"w\"), indent=2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
